{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Figure 6: ablation of the sphere tracing weights\"\"\"\n",
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys\n",
    "import os\n",
    "sys.path.append(os.path.abspath('../'))\n",
    "from common import *\n",
    "from disk_tracing import *\n",
    "\n",
    "import mitsuba as mi\n",
    "\n",
    "from sdf2d.shapes import DiskSDF, UnionSDF\n",
    "\n",
    "# dr.set_log_level(3)\n",
    "fig_name = 'sphere_tracing_weights_ablation'\n",
    "\n",
    "sdf = DiskSDF(mi.Vector2f(0.5, 0.24), mi.Float(0.25))\n",
    "disk2 = DiskSDF(mi.Vector2f(0.3, 0.8), mi.Float(0.15))\n",
    "sdf = UnionSDF(sdf, disk2, k=64, smooth=True)\n",
    "\n",
    "RAY_COLOR = [0.2, 0.3, 0.5]\n",
    "WARP_T_COLOR = [0.8, 0.2, 0.2]\n",
    "SURFACE_COLOR = [1.0, 0.9, 0.7]\n",
    "BG_COLOR = [0.95, 0.95, 0.95]\n",
    "\n",
    "draw_sdf = False\n",
    "base_size = 4\n",
    "n_cols = 2\n",
    "n_rows = 1\n",
    "n_isolines = 31\n",
    "y_offset = -0.15\n",
    "fontsize = 14\n",
    "fig = plt.figure(1, figsize=(n_cols * base_size, n_rows * base_size * 1.02), constrained_layout=False)\n",
    "gs = fig.add_gridspec(n_rows, n_cols, wspace=0.025, hspace=0.025)\n",
    "circle_pos = -0.05\n",
    "ray_res = 1024\n",
    "res = 512\n",
    "params = [False, True]\n",
    "\n",
    "labels = [r\"(a) $\\mathbf{x} + t \\bm{\\omega}$ without $w_{\\mathrm{dist}}$\", r\"(b) $\\mathbf{x} + t \\bm{\\omega}$ with $w_{\\mathrm{dist}}$\"]\n",
    "for col, (label, use_approach_weighting) in enumerate(zip(labels, params)):\n",
    "    ax = fig.add_subplot(gs[col])\n",
    "\n",
    "    ray_origin = mi.Point2f(dr.sin(circle_pos * 2 * dr.pi), dr.cos(circle_pos * 2 * dr.pi)) * (sdf.sdf1.r + 1e-4) + sdf.sdf1.p\n",
    "    n = dr.normalize(sdf.eval_grad(ray_origin))\n",
    "    t = dr.arange(mi.Float, ray_res) / ray_res * dr.pi\n",
    "    T = mi.Matrix2f([[-n.y, n.x],\n",
    "                      [n.x, n.y]])\n",
    "    d = dr.normalize(mi.Vector2f(dr.cos(t), dr.sin(t)))\n",
    "    d = T @ d\n",
    "    return_points = dr.width(t) == 1\n",
    "    dr.set_flag(dr.JitFlag.LoopRecord, not return_points)\n",
    "    ray = mi.Ray2f(ray_origin, d, 0.0, [])\n",
    "    t, warp_t, points, dists, weight_integral = intersect_sdf_simple(sdf, ray, symbolic=True, use_approach_weighting=use_approach_weighting)\n",
    "\n",
    "    y, x = dr.meshgrid(*[dr.linspace(mi.Float, 0, 1, res) for i in range(2)], indexing='ij')\n",
    "    p = mi.Vector2f(x, y)\n",
    "    sdf_values = sdf.eval(p)\n",
    "    r = 0.6\n",
    "    if draw_sdf:\n",
    "        ax.imshow(np.reshape(sdf_values, (res, res)), interpolation='none', extent=[0,1,0,1],\n",
    "                  cmap='coolwarm', vmin=-r, vmax=r, origin='lower')\n",
    "    else:\n",
    "        colors = dr.select(sdf_values < 0, mi.Vector3f(SURFACE_COLOR), mi.Vector3f(BG_COLOR))\n",
    "        ax.imshow(np.reshape(colors, (res, res, 3)), interpolation='none', extent=[0,1,0,1], origin='lower')\n",
    "      \n",
    "    ax.contour(np.reshape(x, (res, res)), np.reshape(y, (res, res)), np.reshape(sdf_values, (res, res)), levels=[0], colors='k')\n",
    "    warp_p = ray(warp_t)\n",
    "    warp_weight = dr.maximum(1 - dr.abs(sdf.eval(warp_p)) / 0.1, 0.0)\n",
    "\n",
    "    # Multiply by weight integral to capture degenerate case of weights being zero\n",
    "    warp_weight *= dr.clamp(weight_integral, 0.0, 1.0)\n",
    "    if col == 0:\n",
    "        warp_weight = np.ones_like(warp_weight)\n",
    "\n",
    "    point_color = np.zeros((ray_res, 4))\n",
    "    point_color[..., 0] = WARP_T_COLOR[0]\n",
    "    point_color[..., 1] = WARP_T_COLOR[1]\n",
    "    point_color[..., 2] = WARP_T_COLOR[2]\n",
    "    point_color[..., -1] = warp_weight\n",
    "\n",
    "    if draw_sdf:\n",
    "        ax.contour(np.reshape(x, (res, res)), np.reshape(y, (res, res)), np.reshape(sdf_values, (res, res)), \n",
    "                        levels=np.linspace(-1,1,n_isolines), alpha=0.9, colors='k', linewidths=0.2)\n",
    "    ax.scatter(warp_p.x, warp_p.y, marker='.', color=point_color, edgecolor='none', s=50, zorder=11)\n",
    "    ax.scatter(ray_origin.x, ray_origin.y, marker='.', color=\"white\", edgecolor=[0.2,0.3,0.6], s=300, zorder=20, lw=1.5)\n",
    "\n",
    "    # Draw a subset of rays \n",
    "    freq = 50\n",
    "    ray_o = np.array(ray.o)[::freq, :] * np.ones((ray_res, 2))[::freq, :]\n",
    "    ray_d = np.array(warp_p - ray.o)[::freq, :]\n",
    "\n",
    "    arrow_color = np.zeros((ray_res, 4))[::freq, :]\n",
    "    arrow_color[..., 0] = RAY_COLOR[0]\n",
    "    arrow_color[..., 1] = RAY_COLOR[1]\n",
    "    arrow_color[..., 2] = RAY_COLOR[2]\n",
    "    arrow_color[..., -1] = 1\n",
    "\n",
    "    ax.quiver(ray_o[:, 0], ray_o[:, 1], ray_d[:, 0], ray_d[:, 1], scale=1, scale_units='xy', color=arrow_color, zorder=15)\n",
    "\n",
    "    ax.set_title(label, y=y_offset, fontsize=fontsize)\n",
    "    ax.set_ylim(0.43, 0.85)\n",
    "    ax.set_xlim(0.1, 0.7)\n",
    "    disable_ticks(ax)\n",
    "\n",
    "    ax.text(ray_origin.x[0] + 0.02, ray_origin.y[0] - 0.02, \"$\\\\mathbf{x}$\", ha='left')\n",
    "\n",
    "    if col == 0:\n",
    "        from matplotlib.lines import Line2D\n",
    "        ax.text(ray_origin.x[0] + 0.1, 0.73, r\"$\\mathbf{x} + t \\bm{\\omega}$\", ha='left')\n",
    "        ax.add_line(Line2D([0.48, 0.51],[0.72,0.73], color='k'))\n",
    "\n",
    "plt.margins(0, 0)\n",
    "# save_fig(fig_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 ('mi')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c8cc192903caa17681aa39d71202092ac11a526d37a1c4ad2948f13605924304"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
